import os
import numpy as np
import matplotlib.pyplot as plt


def read_data_from_txt(file_path):
    y_coords = []
    confidences = []
    with open(file_path, "r") as file:
        for line in file:
            parts = line.strip().split()
            if len(parts) >= 6:
                y_center = float(parts[3])
                if y_center < 0.07 and y_center >0.03:
                    confidence = float(parts[5])
                    y_coords.append(y_center)
                    confidences.append(confidence)
    return y_coords, confidences


def collect_data_from_folder(folder_path):
    indexi = 0
    min_i = 0
    max_i = 10000
    all_y_coords = []
    all_confidences = []
    for file_name in os.listdir(folder_path):
        if file_name.endswith(".txt"):
            indexi = indexi + 1
            if indexi > min_i and indexi < max_i:
                file_path = os.path.join(folder_path, file_name)
                y_coords, confidences = read_data_from_txt(file_path)
                all_y_coords.extend(y_coords)
                all_confidences.extend(confidences)
    return all_y_coords, all_confidences


def calculate_average_confidence(y_coords, confidences, bins):
    bin_indices = np.digitize(y_coords, bins)
    bin_means = [
        np.mean(
            [confidences[i] for i in range(len(confidences)) if bin_indices[i] == j]
        )
        for j in range(1, len(bins))
    ]
    return bin_means


# 文件夹路径
model1_folder = "line_enhanced/labels"
model2_folder = "line_og/labels"

# 区间划分
bins = np.linspace(0, 0.1, 8)  # 将纵坐标划分为20个区间

# 收集数据
model1_y_coords, model1_confidences = collect_data_from_folder(model1_folder)
model2_y_coords, model2_confidences = collect_data_from_folder(model2_folder)
print(min(model1_y_coords))
print(min(model2_y_coords))
print(len(model1_y_coords))
print(len(model2_y_coords))
# 计算每个区间的平均信度
model1_bin_means = calculate_average_confidence(
    model1_y_coords, model1_confidences, bins
)
model2_bin_means = calculate_average_confidence(
    model2_y_coords, model2_confidences, bins
)

# 绘制立方图
plt.figure(figsize=(10, 6))
bin_centers = (bins[:-1] + bins[1:]) / 2
width = (bins[1] - bins[0]) * 0.4

plt.bar(
    bin_centers - width / 2,
    model1_bin_means,
    width=width,
    label="TCone-YOLO",
    alpha=0.5,
)
plt.bar(
    bin_centers + width / 2, model2_bin_means, width=width, label="YOLOv5s", alpha=0.5
)

plt.xlabel("Y Coordinate Interval")
plt.ylabel("Average Confidence")
plt.legend()
plt.show()
